package com.zzy.mccjt.retelimiter;

import com.zzy.mccjt.LimiterUtil;
import com.zzy.mccjt.retelimiter.exception.RateLimiterException;
import lombok.extern.slf4j.Slf4j;
import org.springframework.core.io.ClassPathResource;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.data.redis.core.script.DefaultRedisScript;
import org.springframework.data.redis.core.script.RedisScript;
import org.springframework.scripting.support.ResourceScriptSource;
import org.springframework.stereotype.Component;

import java.time.Instant;
import java.util.Arrays;
import java.util.Collection;
import java.util.HashSet;
import java.util.List;
import java.util.Set;
import java.util.function.Supplier;

/**
 * @program: mccjt-cloud
 * @description: 限流模板
 * @author: zhouyang.zheng
 * @create: 2022-08-01 16:43
 **/
@Slf4j
//@Component
public class RateLimiterTemplate {


    private final RedisTemplate redisTemplate1;

    private final RateLimiterProperties rateLimiterProperties;

    private final RedisScript redisRequestRateLimiterScript;

    public RateLimiterTemplate(RedisTemplate redisTemplate, RateLimiterProperties rateLimiterProperties) {
        this.redisTemplate1 = redisTemplate;
        this.rateLimiterProperties = rateLimiterProperties;
        DefaultRedisScript redisScript = new DefaultRedisScript<>();
        redisScript.setScriptSource(new ResourceScriptSource(
                new ClassPathResource("META-INF/request_rate_limiter.lua")));
        redisScript.setResultType(List.class);
        this.redisRequestRateLimiterScript = redisScript;
    }

    /**
     * 参考Spring Gateway
     *      *  org.springframework.cloud.gateway.filter.ratelimit.isAllowed()
     *      *  redis + lua脚本令牌桶算法限流
     * @param rateLimitCondition
     * @return
     */
    protected boolean check(RateLimitCondition rateLimitCondition) {
        try {
            List<String> keys = LimiterUtil.getKeys(rateLimiterProperties.getNamespace(), rateLimitCondition.getKey());
            /**
             *每秒生产令牌数量
             * 桶容量（并发数量）
             * 时间戳
             * 每次请求消耗令牌数量
             */
            List<Long> results = (List<Long>) this.redisTemplate1.execute(redisRequestRateLimiterScript, keys, rateLimitCondition.getReplenishRate()
                    , rateLimitCondition.getBurstCapacity(), Instant.now().getEpochSecond(), rateLimitCondition.getRequestedTokens());
            assert results != null;
            if (log.isDebugEnabled()) {
                Long tokensLeft = results.get(1);
                log.debug("tokens left: {}", tokensLeft);
            }
            log.info("tokens left: {}", Arrays.asList(results));
            return results.get(0) == 1L;
        } catch (Exception e) {
            log.error(String.format("访问限流异常,请检查配置, key: %s", rateLimitCondition.getKey()), e);
        }
        return true;
    }

    /**
     * 符合单个限流条件时执行
     * @param rateLimitCondition
     * @param func
     * @param <R>
     * @return
     */
    public <R> R executeWithLimit(RateLimitCondition rateLimitCondition, Supplier<R> func) {
        if (!rateLimiterProperties.getEnabled()) {
            return func.get();
        }
        if (this.check(rateLimitCondition)) {
            return func.get();
        }
        throw new RateLimiterException("访问频率超过限制");
    }

    /**
     * 符合多个限流条件时执行
     * @param rateLimitConditions
     * @param func
     * @param <R>
     * @return
     */
    public <R> R executeWithLimit(Collection<RateLimitCondition> rateLimitConditions, Supplier<R> func) {
        if (!rateLimiterProperties.getEnabled()) {
            return func.get();
        }
        boolean allow = true;
        Set<String> keyCache = new HashSet<>(rateLimitConditions.size());
        for (RateLimitCondition rateLimitCondition : rateLimitConditions) {
            String key = rateLimitCondition.getKey();
            if (!keyCache.contains(key)) {
                keyCache.add(key);
                boolean checkResult = this.check(rateLimitCondition);
                if (!checkResult) {
                    allow = false;
                }
            } else {
                log.warn("访问限流存在重复配置, limit type: KEY, key: {}", key);
            }
        }
        if (allow) {
            return func.get();
        }
        throw new RateLimiterException("访问频率超过限制");
    }

}
